import os

import albumentations as A
import cv2
import matplotlib.pyplot as plt

from training.datasets.transform import IsotropicResize


def visualize_augmentations(images_filenames, images_directory, msr_directory, predicted_masks=None, transform=None,
                            num_col=2):
    cols = (5 if predicted_masks else 4) * num_col
    if len(images_filenames) % num_col == 0:
        rows = len(images_filenames) // num_col
    else:
        rows = len(images_filenames) // num_col + 1
    figure, ax = plt.subplots(nrows=rows, ncols=cols, figsize=(8, 12))
    for i, image_filename in enumerate(images_filenames):
        image = cv2.imread(os.path.join(images_directory, image_filename))
        image1 = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image2 = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        image3 = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
        # mask_filename = '{}_diff.png'.format(image_filename[:-4])
        msr = cv2.imread(os.path.join(msr_directory, image_filename), cv2.IMREAD_GRAYSCALE)

        image_filename_split = image_filename[:-4].split('_')
        assert len(image_filename_split) == 2
        sample = {'image': image1, 'image2': image2, 'image3': image3, 'mask': msr}
        transformed = transform(**sample)
        image1 = transformed['image']
        image2 = transformed['image2']
        image3 = transformed['image3']
        msr = transformed['mask']

        # ========== show image ==========
        x_coord = i // num_col
        y_coord = i % num_col * (cols // num_col)
        ax[x_coord, y_coord].imshow(image1)
        ax[x_coord, y_coord + 1].imshow(image2)
        ax[x_coord, y_coord + 2].imshow(image3)
        ax[x_coord, y_coord + 3].imshow(msr, cmap=plt.get_cmap('gray'))

        # ax[x_coord, y_coord].set_title(image_filename)
        # ax[x_coord, y_coord + 1].set_title(mask_filename)
        ax[x_coord, y_coord].set_title(' ')
        ax[x_coord, y_coord + 1].set_title(' ')
        ax[x_coord, y_coord + 2].set_title(' ')
        ax[x_coord, y_coord + 3].set_title(' ')

        ax[x_coord, y_coord].set_axis_off()
        ax[x_coord, y_coord + 1].set_axis_off()
        ax[x_coord, y_coord + 2].set_axis_off()
        ax[x_coord, y_coord + 3].set_axis_off()

        if predicted_masks:
            predicted_mask = predicted_masks[i]
            ax[x_coord, y_coord + 2].imshow(predicted_mask, interpolation="nearest")
            ax[x_coord, y_coord + 2].set_title("Predicted mask")
            ax[x_coord, y_coord + 2].set_axis_off()
    plt.tight_layout()
    plt.show()


def main():
    root_dir = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU'
    crops = os.path.join(root_dir, 'Test_files_crops')
    msr = os.path.join(root_dir, 'Test_files_MSRCR')
    # landmarks = os.path.join(root_dir, 'landmarks')

    video_name = '6_3_48_1'
    images_path = os.path.join(crops, video_name)
    msr_path = os.path.join(msr, video_name)
    # video_name_split = video_name.split('_')
    # landmarks_path = os.path.join(landmarks, '{}_{}'.format(video_name_split[0], video_name_split[2]))

    frames = os.listdir(images_path)
    # len(frames)

    input_size = (3, 224, 224)
    # resize = int(input_size[1] / 0.875)
    resize = int(input_size[1] / 1)
    train_transform = A.Compose([
        A.HorizontalFlip(),
        A.OneOf([
            IsotropicResize(max_side=resize, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),
            IsotropicResize(max_side=resize, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_LINEAR),
            IsotropicResize(max_side=resize, interpolation_down=cv2.INTER_LINEAR, interpolation_up=cv2.INTER_LINEAR),
        ], p=1),
        A.PadIfNeeded(min_height=resize, min_width=resize, border_mode=cv2.BORDER_CONSTANT, value=0),
        # A.RandomCrop(height=input_size[1], width=input_size[2]),
        # A.OneOf([A.RandomBrightnessContrast(), A.FancyPCA(), A.HueSaturationValue()], p=0.3),
        # A.ToGray(p=0.2),
        # A.ShiftScaleRotate(rotate_limit=(-20, 20), border_mode=cv2.BORDER_CONSTANT, p=0.3),
    ], additional_targets={'image2': 'image', 'image3': 'image'})

    generalization_transform = A.Compose([
        A.Blur(blur_limit=(5, 20), p=0.5),
        A.OneOf([A.RandomBrightnessContrast(), A.FancyPCA(), A.HueSaturationValue()], p=0.7),
        # A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, border_mode=cv2.BORDER_CONSTANT, p=0.5),
        # A.OpticalDistortion(distort_limit=(1., 2.), border_mode=cv2.BORDER_CONSTANT, p=0.5)
    ])

    visualize_augmentations(frames[:16], images_path, msr_path, transform=train_transform)


if __name__ == '__main__':
    main()
